"""
鲸鱼优化算法（WOA）调度器
实现基于鲸鱼优化算法的单目标优化调度
"""

import numpy as np
from typing import List, Dict, Any, Tuple
import random
from datetime import datetime, timedelta
from .base_scheduler import BaseScheduler

class WOAScheduler(BaseScheduler):
    """鲸鱼优化算法调度器"""
    
    def __init__(self, config: Dict[str, Any] = None):
        super().__init__(config)
        self.max_iterations = self.config.get('max_iterations', 100)
        self.population_size = self.config.get('population_size', 30)
        self.bounds = self.config.get('bounds', {'lower': -1, 'upper': 1})
        
    def optimize(self, 
                 resources: List[Dict], 
                 appointments: List[Dict], 
                 objectives: List[str]) -> Dict[str, Any]:
        """
        执行WOA优化
        
        Args:
            resources: 资源列表
            appointments: 预约列表
            objectives: 优化目标列表
            
        Returns:
            调度结果字典
        """
        self.resources = resources
        self.appointments = appointments
        
        # 初始化鲸鱼种群
        whales = self._initialize_whales()
        
        # 找到初始最优解
        best_whale = self._find_best_whale(whales)
        best_fitness = self._evaluate_whale(best_whale)
        
        # WOA主循环
        for iteration in range(self.max_iterations):
            a = 2 - iteration * (2 / self.max_iterations)  # 线性递减参数a
            
            for i, whale in enumerate(whales):
                r1, r2 = random.random(), random.random()
                A = 2 * a * r1 - a
                C = 2 * r2
                b = 1  # 螺旋形状参数
                l = (a - 1) * random.random() + 1
                
                p = random.random()
                
                if p < 0.5:
                    if abs(A) >= 1:
                        # 搜索猎物（探索）
                        rand_leader_index = random.randint(0, len(whales) - 1)
                        X_rand = whales[rand_leader_index]
                        D_X_rand = abs(C * X_rand - whale)
                        whale = X_rand - A * D_X_rand
                    else:
                        # 包围猎物（开发）
                        D_Leader = abs(C * best_whale - whale)
                        whale = best_whale - A * D_Leader
                else:
                    # 螺旋更新位置
                    distance2Leader = abs(best_whale - whale)
                    whale = distance2Leader * np.exp(b * l) * np.cos(l * 2 * np.pi) + best_whale
                
                # 边界处理
                whale = self._apply_bounds(whale)
                whales[i] = whale
                
                # 更新最优解
                fitness = self._evaluate_whale(whale)
                if fitness < best_fitness:
                    best_fitness = fitness
                    best_whale = whale.copy()
        
        # 将最优解转换为调度方案
        solution = self._whale_to_solution(best_whale)
        objectives_values = self.calculate_objectives(solution)
        
        return {
            'solution': solution,
            'objectives': objectives_values,
            'algorithm': 'WOA',
            'iterations': self.max_iterations,
            'best_fitness': best_fitness
        }
    
    def _initialize_whales(self) -> List[np.ndarray]:
        """初始化鲸鱼种群"""
        whales = []
        dimension = len(self.appointments) * 2  # 每个预约有时间和资源两个维度
        
        for _ in range(self.population_size):
            whale = np.random.uniform(
                self.bounds['lower'], 
                self.bounds['upper'], 
                dimension
            )
            whales.append(whale)
            
        return whales
    
    def _find_best_whale(self, whales: List[np.ndarray]) -> np.ndarray:
        """找到最优鲸鱼"""
        best_whale = whales[0]
        best_fitness = self._evaluate_whale(best_whale)
        
        for whale in whales[1:]:
            fitness = self._evaluate_whale(whale)
            if fitness < best_fitness:
                best_fitness = fitness
                best_whale = whale
                
        return best_whale
    
    def _evaluate_whale(self, whale: np.ndarray) -> float:
        """评估鲸鱼适应度"""
        # 将鲸鱼位置转换为调度方案
        solution = self._whale_to_solution(whale)
        
        # 计算综合目标函数
        objectives = self.calculate_objectives(solution)
        
        # 加权组合多个目标
        weights = [0.3, 0.3, 0.2, 0.2]  # 资源利用率、等待时间、公平性、加班成本
        fitness = (
            weights[0] * (1 - objectives['resource_utilization']) +  # 最大化资源利用率
            weights[1] * objectives['patient_wait_time'] +          # 最小化等待时间
            weights[2] * (1 - objectives['fairness']) +             # 最大化公平性
            weights[3] * objectives['overtime_cost'] / 1000        # 最小化加班成本
        )
        
        return fitness
    
    def _whale_to_solution(self, whale: np.ndarray) -> Dict[str, Any]:
        """将鲸鱼位置转换为调度方案"""
        assignments = {}
        appointment_count = len(self.appointments)
        
        for i, appointment in enumerate(self.appointments):
            # 从鲸鱼位置解码资源分配
            resource_index = int((whale[i] + 1) / 2 * len(self.resources)) % len(self.resources)
            resource = self.resources[resource_index]
            
            # 从鲸鱼位置解码时间分配
            time_offset = whale[appointment_count + i] * 60  # 时间偏移（分钟）
            start_time = appointment['preferred_time'] + timedelta(minutes=time_offset)
            duration = appointment.get('duration', 30)
            end_time = start_time + timedelta(minutes=duration)
            
            resource_id = resource['id']
            if resource_id not in assignments:
                assignments[resource_id] = []
                
            assignments[resource_id].append({
                'appointment_id': appointment['id'],
                'resource_id': resource_id,
                'start_time': start_time,
                'end_time': end_time
            })
            
        return {'assignments': assignments}
    
    def _apply_bounds(self, whale: np.ndarray) -> np.ndarray:
        """应用边界约束"""
        return np.clip(whale, self.bounds['lower'], self.bounds['upper'])
